import numpy as np
import cv2
from open3d import geometry
from cyw_devkit.core.io import read_points
from cyw_devkit.core.visualizer import zwh_viewer
from PIL import Image


def trans_img_new_intrin(out_intrin, new_size, old_intrin, img):
    H = np.dot(out_intrin, np.linalg.inv(old_intrin))
    img = Image.fromarray(img[:, :, ::-1])

    # img.show()

    w, h = img.size
    H_inv = np.linalg.inv(H)
    H_inv = H_inv / H_inv[2, 2]  # 归一化
    # PIL expects a 3x3 matrix for perspective transform
    coefficients = H_inv.flatten()[:8]  # 取前8个元素，用于透视变换
    # 进行仿射变换
    if isinstance(new_size, list):
        w = new_size[1]
        h = new_size[0]
    new_img = img.transform(
        (w, h),
        Image.PERSPECTIVE,
        coefficients,
        Image.BICUBIC
    )
    # new_img = self.remove_black_borders(new_img)
    # self.show_diff_intrin_imgs(img, new_img)
    # img.show()
    # new_img.show()

    # new_img.save(filename)
    return new_img

if __name__ == '__main__':
    # K= np.array([[1001.5819983664084,
    #     0.0,
    #         950.5118842949007
    #     ],
    #     [
    #         0.0,
    #         1000.9117211003688,
    #         496.8809954299153
    #     ],
    #     [
    #         0.0,
    #         0.0,
    #         1.0
    #     ]
    # ])
    # D=  np.array([
    #     -0.37047944718384757,
    #     0.1674906367136755,
    #     -0.0011702109119313324,
    #     -0.0001258841058397891,
    #     -4.0070333591922905e-05
    # ])

    K= np.array([[1.0169827755366000e+03, 0., 9.5971434761570004e+02],
                  [0.0,1.0169114728239000e+03, 5.3979974781310000e+02],
                  [0.0, 0.0, 1.]])
    D= np.array([-0.37047944718384757, 0.1674906367136755, 2.8021091193133233e-05, -0.0002858841058397891, -0.040070333591922906])
    RT= np.array([[5.1722458881657274e-03, -9.9992400037234830e-01,-1.1191128264489036e-02, 1.2805628752804571e-01],
        [-8.3640057748189944e-03,1.1147627927642734e-02, -9.9990288218355761e-01,5.0922666642886968e-03],
        [9.9995164447080698e-01, 5.2653462323700162e-03,-8.3057118407829922e-03,-3.8021243351345185e-02],
         [0,0,0,1]])

    lidar2base = np.array([[ 1.0, 0.0, 0.0,1.08],
                           [ 0.0,  1.0,  0.0,-0.845],
                           [ 0.0,  0.0,  1.0,1.087],
                           [0,0,0,1 ]])

    radar2lidar = np.array([
        [0.999956, -0.00941348, -5.12898e-05, -0.0249084],
         [0.00941242, 0.999901, -0.0105006, 0.240744],
         [0.000150132, 0.0104997, 0.999945, -0.0110709],
         [0, 0, 0, 1]
    ])
    radar2base = lidar2base @ radar2lidar

    points = read_points('/media/adt/ZWH4T/ZWH/Dataset3d/dataset_18xx-bag/cyw2.0/tmx/datasets/2025-04-02-10-28-13_0_filter0_follow/lidar/pointcloud_lidar3/1743590369345.bin',dim=5)
    radar = read_points('/media/adt/ZWH4T/ZWH/Dataset3d/dataset_18xx-bag/cyw2.0/tmx/datasets/2025-04-02-10-28-13_0_filter0_follow/radar/pointcloud_radar3/1743590369352.bin',dim=5)

    radar = zwh_viewer.tranform_pt_pad(radar, tfs=[radar2base])

    points = zwh_viewer.tranform_pt_pad(points, tfs=[lidar2base])

    img = cv2.imread('/media/adt/ZWH4T/ZWH/Dataset3d/dataset_18xx-bag/cyw2.0/tmx/datasets/2025-04-02-10-28-13_0_filter0_follow/camera/camera75/1743590369433.jpg')

    img = cv2.resize(img, (1920, 1080))

    newcameramtx, roi = cv2.getOptimalNewCameraMatrix(K,
                                                      D,
                                                      (1920, 1080),
                                                      0)
    # 畸变校正
    # img = cv2.resize(img, (cam_data['imgw'], cam_data['imgh']))
    dst = cv2.undistort(img, K, D, None, newcameramtx)

    out_intrin = [
        [680.0237639534591,0.0,640.0],
        [0.0,650.0237639534591,360.0],
        [0.0,0.0,1.0]]
    # cv2.getOptimalNewCameraMatrix(xxx,,,0)
    # out_intrin = [
    #     [570.0,0.0,640.0],
    #     [0.0,570.0,360.0],
    #     [0.0,0.0,1.0]]
    new_size = [720, 1280]

    dst = np.array(trans_img_new_intrin(out_intrin, new_size, newcameramtx, dst))

    K_pad = np.eye(4)
    K_pad[:3,:3] = out_intrin

    cam2base =  lidar2base @ np.linalg.inv(RT)
    img2 = zwh_viewer.draw_pts2img(points,dst,tfs=[np.linalg.inv(cam2base),K_pad])

    img2 = zwh_viewer.draw_pts2img(radar,img2,tfs=[np.linalg.inv(cam2base),K_pad], pts_size=4)

    cv2.namedWindow('undistort', cv2.WINDOW_NORMAL)
    cv2.imshow('undistort', img2)
    cv2.waitKey(0)

    pass
